# encoding: utf-8
# author: xincz

from torch.utils.data import Dataset, DataLoader
from torchvision import transforms


class ImageDataset(Dataset):
    def __init__(self, filelist, data_tensor):
        super(ImageDataset, self).__init__()
        self.filelist = filelist
        self.data_tensor = data_tensor
        self.trans = transforms.Compose([
            transforms.Pad(30, fill=0, padding_mode='constant'),
            transforms.CenterCrop((46, 31))
        ])

    def __len__(self):
        return len(self.filelist)

    def __getitem__(self, index):
        name = self.filelist[index]
        image = (self.trans(self.data_tensor[name])/127.5).float()
        return image, name
